#ifndef STREAMMANAGER_H
#define STREAMMANAGER_H

#include <QIODevice>
#include <QDataStream>
#include <QBuffer>
#include <QAbstractSocket>
#include <QJsonObject>
#include <QJsonDocument>
#include <QImage>
#include <QTcpSocket>
#include <QByteArray>
#include "compressedqimage.h"
#include "constants.h"

class BlockReader
{
public:
    BlockReader(QIODevice *io)
    {
        buffer.open(QIODevice::ReadWrite);
        _stream.setVersion(QDataStream::Qt_4_8);
        _stream.setDevice(&buffer);

        quint64 blockSize;

        // Read the size.
        readMax(io, sizeof(blockSize));
        buffer.seek(0);
        _stream >> blockSize;

        // Read the rest of the data.
        readMax(io, blockSize);
        buffer.seek(sizeof(blockSize));
    }

    QDataStream &stream()
    {
        return _stream;
    }

private:
    // Blocking reads data from socket until buffer size becomes exactly n. No
    // additional data is read from the socket.
    void readMax(QIODevice *io, int n)
    {
        while (buffer.size() < n)
        {
            while (io->bytesAvailable() <= 0)
            {
                io->waitForReadyRead(100);
            }
            buffer.write(io->read(n - buffer.size()));
        }
    }
    QBuffer buffer;
    QDataStream _stream;
};

class BlockWriter
{
public:
    BlockWriter(QIODevice *io)
    {
        buffer.open(QIODevice::WriteOnly);
        this->io = io;
        _stream.setVersion(QDataStream::Qt_4_8);
        _stream.setDevice(&buffer);

        // Placeholder for the size. We will get the value
        // at the end.
        _stream << quint64(0);
    }

    ~BlockWriter()
    {
        // Write the real size.
        _stream.device()->seek(0);
        _stream << (quint64)buffer.size();

        // Flush to the device.
        io->write(buffer.buffer());
    }

    QDataStream &stream()
    {
        return _stream;
    }

private:
    QBuffer buffer;
    QDataStream _stream;
    QIODevice *io;
};

inline void receive_map(QIODevice *socket, QVariantMap &map)
{
    QByteArray arr;
    BlockReader(socket).stream() >> arr;
    QJsonDocument doc = QJsonDocument::fromJson(arr);
    QJsonObject obj = doc.object();
    map = obj.toVariantMap();
}

inline void send_map(QTcpSocket *socket, QVariantMap &map)
{
    QJsonObject obj = QJsonObject::fromVariantMap(map);
    QJsonDocument doc = QJsonDocument(obj);
    QByteArray arr = doc.toJson();
    BlockWriter(socket).stream() << arr;
    while (socket->flush());
}

#ifndef USE_H264

#ifdef CLIENT
inline void receive_image(QIODevice *socket, QImage &image)
{
#ifdef QIMAGE_DIFF_COMPRESS
    QVariantMap map;
    receive_map(socket, map);
    QByteArray arr;
    BlockReader(socket).stream() >> arr;
#ifdef QIMAGE_COMPRESS
    arr = qUncompress(arr);
#endif
    if (map.value("image_type").toInt() == FULL_IMAGE)
    {
        image = QImage::fromData(arr);
    }
    else if (map.value("image_type").toInt() == PART_IMAGE)
    {
        image = CompressedQImage::addTo(arr, image);
    }
#else
    QByteArray arr;
    BlockReader(socket).stream() >> arr;
#ifdef QIMAGE_COMPRESS
    arr = qUncompress(arr);
#endif
    image = QImage::fromData(arr);
#endif
}

#endif /* CLIENT */

#ifdef SERVER
inline void send_image(QTcpSocket *socket, QImage &oldimage, QImage &newimage)
{
#ifdef QIMAGE_DIFF_COMPRESS
    CompressedQImage compressedImage(oldimage, newimage);

    QByteArray arr;
    QBuffer buffer(&arr);
    buffer.open(QIODevice::WriteOnly);
    newimage.save(&buffer, QIMAGE_SAVE_FORMAT, QIMAGE_SAVE_QUALITY);
    buffer.close();

    QVariantMap map;

    if (compressedImage.getValue().size() <= 0 || arr.size() <= compressedImage.getValue().size())
    {
        //        info("send image: FULL_IMAGE");
        map.insert("image_type", FULL_IMAGE);
        send_map(socket, map);
#ifdef QIMAGE_COMPRESS
        arr = qCompress(arr, 9);
#endif
        BlockWriter(socket).stream() << arr;
    }
    else
    {
        //        info("send image: PART_IMAGE");
        map.insert("image_type", PART_IMAGE);
        send_map(socket, map);
        arr = compressedImage.getValue();
#ifdef QIMAGE_COMPRESS
        arr = qCompress(arr, 9);
#endif
        BlockWriter(socket).stream() << arr;
    }
#else
    QByteArray arr;
    QBuffer buffer(&arr);
    buffer.open(QIODevice::WriteOnly);
    newimage.save(&buffer, QIMAGE_SAVE_FORMAT, QIMAGE_SAVE_QUALITY);
    buffer.close();

#ifdef QIMAGE_COMPRESS
    arr = qCompress(arr, 9);
#endif
    BlockWriter(socket).stream() << arr;
#endif

    while (socket->flush());
}

#endif /* SERVER */

#else /* USE_H264 */


#ifdef CLIENT
#include "decoder.h"
static uchar result[1920 * 1080 * 4];

inline void receive_image(QIODevice *socket, QImage &image, Decoder *decoder)
{
    QByteArray array;
    BlockReader(socket).stream() >> array;
#ifdef H264_COMPRESS
    array = qUncompress(array);
#endif
    int out_size;

    int width, height;
    int err = decoder->loop((uint8_t *)array.data(), array.size(), result, &out_size, &width, &height);

    if (err == 0){
//        info("out_size = %d", out_size);
        info("width = %d height = %d", width, height);
        image = QImage(result, width, height, QImage::Format_RGB888);
    }
    else {
        info("err = %d", err);
    }
}
#endif /* CLIENT */

#ifdef SERVER
#include "encoder.h"
#include "flv2x264converter.h"
static int pts = 0;

inline void send_image(QTcpSocket *socket, Encoder *encoder, QImage &newimage, bool & sent_header)
{
    //Header
    QByteArray inbits, outbits;
    if (!sent_header){
        info("send header");
        inbits = encoder->getSequenceHeader();
        to_h264(inbits.data(), inbits.size(), outbits);
        BlockWriter(socket).stream() << outbits;
        sent_header = true;
    }

    //Data
    inbits = encoder->getEncodedFrame(newimage, pts);
    if (inbits.size() > 0){
        info("send data");
        to_h264(inbits.data(), inbits.size(), outbits);
        BlockWriter(socket).stream() << outbits;
    }
    pts += 100;

    while (socket->flush());
}
#endif /* SERVER */


#endif /* USE_H264 */

#endif
